#include "rtp.h"
#include <string.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>

static int h264_rtp_pack(struct rtp_packetizer *packetizer, char *buffer, int len, uint32_t timestamp)
{
    if (!packetizer) return -1;
    struct rtp_packet *packet = packetizer->packet;
    uint16_t *seq = &packetizer->seq;
    if (buffer[0] == 0x00 && buffer[1] == 0x00 && buffer[2] == 0x01)
    {
        buffer += 3;
        len -= 3;
    }
    else if (buffer[0] == 0x00 && buffer[1] == 0x00 && buffer[2] == 0x00 && buffer[3] == 0x01)
    {
        buffer += 4;
        len -= 4;
    }
    int ret = 0;
    if (len <= MAX_RTP_PAYLOAD_SIZE)
    {
        packet->header.marker = 1;
        packet->header.timestamp = htonl(timestamp);
        packet->header.seq = htons(*seq);
        memcpy(packet->payload, buffer, len);
        packetizer->callback(packet, len + 12, packetizer->usr_data);
        *seq += 1;
    }
    else
    {
        int pos = 0;
        int nalu_type = buffer[0];
        buffer += 1;
        len -= 1;
        packet->header.payloadType = RTP_PAYLOAD_TYPE_H26X;
        packet->header.timestamp = htonl(timestamp);
        int pk_count = len / MAX_RTP_PAYLOAD_SIZE;
        int remain_size = len % MAX_RTP_PAYLOAD_SIZE;
        for (int i = 0; i < pk_count; i++)
        {
            memset(packet->payload, 0, 2);
            packet->header.marker = 0;
            packet->header.seq = htons(*seq);
            packet->payload[0] = (nalu_type & 0xE0) | 28;
            packet->payload[1] = nalu_type & 0x1F;
            if (i == 0)
            {
                packet->payload[1] |= 0x80;
            }
            else if (remain_size == 0 && i == pk_count - 1)
            {
                packet->payload[1] |= 0x40;
            }
            memcpy(packet->payload + 2, buffer + pos, MAX_RTP_PAYLOAD_SIZE);
            packetizer->callback(packet, MAX_RTP_PAYLOAD_SIZE + 12 + 2, packetizer->usr_data);
            pos += MAX_RTP_PAYLOAD_SIZE;
            *seq += 1;
        }

        if (remain_size > 0)
        {
            packet->header.marker = 1;
            packet->header.seq = htons(*seq);
            packet->payload[0] = (nalu_type & 0xE0) | 28;
            packet->payload[1] = nalu_type & 0x1F;
            packet->payload[1] |= 0x40;
            memcpy(packet->payload + 2, buffer + pos, remain_size);
            packetizer->callback(packet, remain_size + 12 + 2, packetizer->usr_data);
            *seq += 1;
        }
    }
    return ret;
}

static int h265_rtp_pack(struct rtp_packetizer *packetizer, char *buffer, int len, uint32_t timestamp)
{
    if (!packetizer) return -1;
    struct rtp_packet *packet = packetizer->packet;
    uint16_t *seq = &packetizer->seq;
    if (buffer[0] == 0x00 && buffer[1] == 0x00 && buffer[2] == 0x01)
    {
        buffer += 3;
        len -= 3;
    }
    else if (buffer[0] == 0x00 && buffer[1] == 0x00 && buffer[2] == 0x00 && buffer[3] == 0x01)
    {
        buffer += 4;
        len -= 4;
    }

    int ret = 0;
    if (len <= MAX_RTP_PAYLOAD_SIZE)
    {
        packet->header.marker = 1;
        packet->header.timestamp = htonl(timestamp);
        packet->header.seq = htons(*seq);
        memcpy(packet->payload, buffer, len);
        packetizer->callback(packet, len + 12, packetizer->usr_data);
        *seq += 1;
    }
    else
    {
        int pos = 2;
        int nalu_type = (buffer[0] & 0x7F) >> 1;
        packet->header.payloadType = RTP_PAYLOAD_TYPE_H26X;
        packet->header.timestamp = htonl(timestamp);
        int pk_count = len / MAX_RTP_PAYLOAD_SIZE;
        int remain_size = len % MAX_RTP_PAYLOAD_SIZE;
        for (int i = 0; i < pk_count; i++)
        {
            memset(packet->payload, 0, 3);
            packet->header.marker = 0;
            packet->header.seq = htons(*seq);
            packet->payload[0] = (buffer[0] & 0x81) | (49 << 1);
            packet->payload[1] = buffer[1];
            packet->payload[2] = nalu_type;
            if (i == 0)
            {
                packet->payload[2] |= 0x80;
            }
            else if (remain_size == 0 && i == pk_count - 1)
            {
                packet->header.marker = 1;
                packet->payload[2] |= 0x40;
            }
            memcpy(packet->payload + 3, buffer + pos, MAX_RTP_PAYLOAD_SIZE);
            packetizer->callback(packet, MAX_RTP_PAYLOAD_SIZE + 12 + 3, packetizer->usr_data);
            pos += MAX_RTP_PAYLOAD_SIZE;
            *seq += 1;
        }

        if (remain_size > 0)
        {
            packet->header.marker = 1;
            packet->header.seq = htons(*seq);
            memset(packet->payload, 0, 3);
            packet->payload[0] = (buffer[0] & 0x81) | (49 << 1);
            packet->payload[1] = buffer[1];
            packet->payload[2] = nalu_type;
            packet->payload[2] |= 0x40;
            memcpy(packet->payload + 3, buffer + pos, remain_size);
            packetizer->callback(packet, remain_size + 12 + 3, packetizer->usr_data);
            *seq += 1;
        }
    }
    return ret;
}

static int aac_rtp_pack(struct rtp_packetizer *packetizer, char *buffer, int len, uint32_t timestamp)
{
    if (!packetizer) return -1;
    struct rtp_packet *packet = packetizer->packet;
    uint16_t *seq = &packetizer->seq;
    char AU[4] = {0};
    AU[0] = 0x00;
    AU[1] = 0x10;
    AU[2] = (len & 0x1fe0) >> 5;
    AU[3] = (len & 0x1f) << 3;

    packet->header.marker = 1;
    packet->header.timestamp = htonl(timestamp);
    packet->header.seq = htons(*seq);
    *seq += 1;

    memcpy(packet->payload, AU, 4);
    memcpy(packet->payload + 4, buffer, len);
    packetizer->callback(packet, len + 12 + 4, packetizer->usr_data);

    return len + 12 + 4;
}

static int pcma_rtp_pack(struct rtp_packetizer *packetizer, char *buffer, int len, uint32_t timestamp)
{
    if (!packetizer) return -1;
    struct rtp_packet *packet = packetizer->packet;
    uint16_t *seq = &packetizer->seq;
    packet->header.marker = 1;
    packet->header.timestamp = htonl(timestamp);
    packet->header.seq = htons(*seq);
    *seq += 1;

    memcpy(packet->payload, buffer, len);
    packetizer->callback(packet, len + 12, packetizer->usr_data);

    return len + 12;
}

struct rtp_packetizer* create_rtp_packetizer(int payloadType, uint32_t ssrc, rtp_data_callback callback, void* usr_data)
{
    if (!callback) return NULL;
    struct rtp_packetizer* packetizer = (struct rtp_packetizer*)malloc(sizeof(struct rtp_packetizer));
    if (!packetizer)
        return NULL;

    int payload = RTP_PAYLOAD_TYPE_H26X;
    memset(packetizer, 0, sizeof(struct rtp_packetizer));
    switch (payloadType)
    {
    case RTP_H264:
        packetizer->input = h264_rtp_pack;
        payload = RTP_PAYLOAD_TYPE_H26X;
        break;
    case RTP_H265:
        packetizer->input = h265_rtp_pack;
        payload = RTP_PAYLOAD_TYPE_H26X;
        break;
    case RTP_AAC:
        packetizer->input = aac_rtp_pack;
        payload = RTP_PAYLOAD_TYPE_AAC;
        break;
    case RTP_PCMA:
        packetizer->input = pcma_rtp_pack;
        payload = RTP_PAYLOAD_TYPE_PCMA;
        break;
    default:
        free(packetizer);
        return NULL;
    }
    packetizer->callback = callback;
    packetizer->payloadType = payloadType;
    packetizer->usr_data = usr_data;
    packetizer->packet = (struct rtp_packet*)malloc(MAX_RTP_PAYLOAD_SIZE);
    if (!packetizer->packet) {
        free(packetizer);
        return NULL;
    }
    packetizer->packet->header.payloadType = payload;
    packetizer->packet->header.ssrc = htonl(ssrc);
    packetizer->packet->header.csrcLen = 0;
    packetizer->packet->header.extension = 0;
    packetizer->packet->header.padding = 0;
    packetizer->packet->header.version = RTP_VESION;

    return packetizer;
}

int rtp_payload_encode_input(struct rtp_packetizer *packetizer, void *data, int len, uint32_t timestamp)
{
    if (!packetizer) return -1;
    return packetizer->input(packetizer, data, len, timestamp);
}

void rtp_packetizer_destroy(struct rtp_packetizer *packetizer)
{
    if (!packetizer) return;
    if (!packetizer->packet) free(packetizer->packet);
    free(packetizer);
}
